package textprocessing.nlpanalysis.context_annotate;

import algorithm.config.HumanLanguage;
import algorithm.nlp.corenlp.TextAnnotator;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import edu.stanford.nlp.coref.CorefCoreAnnotations;
import edu.stanford.nlp.coref.data.CorefChain;
import edu.stanford.nlp.ling.CoreAnnotations;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.ling.Label;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.util.CoreMap;
import edu.stanford.nlp.util.SystemUtils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.tools.cmd.gen.AnyVals;
import semantic_compute.entity_linking.schema.EntityLinkingResult;
import textprocessing.nlpanalysis.constituent.ContextGenerator;
import textprocessing.nlpanalysis.constituent.ContextItem;
import textprocessing.nlpanalysis.entity_linking.TextEntityLinker;

import java.util.*;

/**
 * Created by julianzliu on 4/19/2017.
 */

/*******************************
 *
 * 输入：一个段落，包含多个句子
 * 输出：标注后的 context 集合
 *
 */
public class ContextAnnotation {

    static final Logger logger = LoggerFactory.getLogger(ContextAnnotation.class);

    private TextAnnotator textAnnotator;
    private ContextGenerator contextGenerator;
    private TextEntityLinker textEntityLinker;
    private Map<Integer,Map<Integer,Set<String>>> sentenceEndtokenEntitySet;

    private ObjectMapper mapper;

    public ContextAnnotation(){

        this.textAnnotator = new TextAnnotator(HumanLanguage.ENGLISH);
        this.contextGenerator = new ContextGenerator();
        this.textEntityLinker = new TextEntityLinker();
        this.sentenceEndtokenEntitySet = new HashMap<>();

        this.mapper = new ObjectMapper();
    }


    /**************************************
     * 入口函数
     * 输入文本，返回标注结果
     * @param text
     * @return
     */
    public List<ContextResult> getContextsOfDocument(String text){
        List<ContextResult> allContexts = new ArrayList<>();

        try{
             /*
            1) 自然语言处理
             */
            Annotation document = this.textAnnotator.getAnnotatedText( text );

            /*
            2）实体链接，并将相同共指消解链上的词语指向相同  实体
             */

            this.sentenceEndtokenEntitySet = getSentenceEndtokenEntitySet(text, document);

            try{

                //this.sentenceEndtokenEntitySet = getSentenceEndtokenEntitySet(text, document);

            }catch (Exception e){
                logger.error("[getContextsOfDocument] Error {}", e.getMessage());
            }



            /*
            3) 基本成分识别
             */
            List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);
            for(CoreMap sentence: sentences){
                //List<ContextItem> senContextItems = this.contextGenerator.generateContextsOfSentence(sentence);
                List<ContextResult> senContetxts = getContextsOfSentence( sentence, document );
                allContexts.addAll(senContetxts);
            }
            //

        }catch (Exception e){
            logger.error("[getContextsOfDocument] Error Gen Contetxs! {}", e.getMessage());
        }


        return allContexts;
    }


    public ObjectNode getContextsOfDocumentsJson(String text){
        ObjectNode docContexts = this.mapper.createObjectNode();

        return  docContexts;
    }



    public List<ContextResult> getContextsOfSentence(CoreMap sentence, Annotation document){
        List<ContextResult> senContexts = new ArrayList<>();
        List<CoreLabel> senTokens = sentence.get(CoreAnnotations.TokensAnnotation.class);
        //Tree senTree = sentence.get(TreeCoreAnnotations.TreeAnnotation.class);
        List<ContextItem> senContextItems = this.contextGenerator.generateContextsOfSentence(sentence);
        for(ContextItem item: senContextItems){
            ContextResult annotateSingleContext  = annotateSingleContext( item, senTokens, document );
            senContexts.add( annotateSingleContext );
        }
        return senContexts;
    }


    /*
    将Context中出现的实体  进行 标注
     */
    public ContextResult annotateSingleContext(ContextItem item, List<CoreLabel> senTokens, Annotation document){
        ContextResult context = new ContextResult();
        List<Tree> treeTokens = item.getTokens();
        System.out.println("\n\n[annotateSingleContext] try to get entity linking and coref resolve\n\n");

        for(Tree treeToken: treeTokens){
            try{
                if(!treeToken.isLeaf()){
                    logger.error("[annotateSingleContext] contian None Leaf Tree Node, info = \n{}", treeToken.toString());
                    throw new Exception("None Leaf Tree Node Error");
                }else{

                    Integer tokenIndex = getTokenIndexOfLeafNode(treeToken); //该下标从1开始
                    Integer index = tokenIndex - 1; //转换为数组下标
                    CoreLabel matchedToken = senTokens.get(index);
                    context.context += matchedToken.word() + " ";

                    // 指代消解
                    List<CoreLabel> repTokens = getRepresentativeMentionTokensOfaToken( matchedToken, document );
                    for(CoreLabel token: repTokens){
                        context.context += token.word() + " ";
                    }

                    // 实体链接
                    Set<String> entities = getEntitiesOfToken(matchedToken);
                    if(entities != null)
                        context.extendEntities( entities );
                }
            }
            catch (Exception e){
                e.printStackTrace();
            }

        }
        System.out.println("[annotateSingleContext] finished a Context!\n\n");
        return context;
    }

    private Integer getTokenIndexOfLeafNode(Tree token){
        Label label = token.label();
        String[] labelArray = label.toString().split("-");
        String indexStr = labelArray[labelArray.length-1];
        Integer index = Integer.parseInt(indexStr);
        return index;
    }

    public List<CoreLabel> getRepresentativeMentionTokensOfaToken(CoreLabel matchedToken, Annotation document){
        //System.out.println("[]");
        /*
        注意：
        1） 凡是和展示结果先关的下标，均是从1开始
        2） 存储在数组中的实体，下标从0开始
        ---------------------
        1) CorefMention 下标从1开始
        2) CoreLabel 下标从0开始
         */
        Integer corefClustId= matchedToken.get(CorefCoreAnnotations.CorefClusterIdAnnotation.class);
        System.out.println("[matchedToen = " + matchedToken.word() +  " | corefClusterID = " + corefClustId + " ]");
        Map<Integer, CorefChain> corefs = document.get(CorefCoreAnnotations.CorefChainAnnotation.class);
        CorefChain chain = corefs.get(corefClustId);
        System.out.println("[matchedChain = " + chain + "]");

        List<CoreLabel> repTokens = new ArrayList<>();
        if(chain == null){
            //repTokens.add(matchedToken);
        }else{
            CorefChain.CorefMention reprMent = chain.getRepresentativeMention();
            System.out.println(reprMent);

            int sentINdx = reprMent.sentNum -1;

            List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);
            CoreMap corefSentence = sentences.get(sentINdx);
            List<CoreLabel> corefSentenceTokens = corefSentence.get(CoreAnnotations.TokensAnnotation.class);

            //System.out.println("token.sentIndex = " + matchedToken.sentIndex() + " , rep.sentNum = " + sentINdx);
            //System.out.println("token.index = " + matchedToken.index() + " , rep.start = " + reprMent.startIndex + " , rep.end = " + reprMent.endIndex);
            // 不是rep自身
            if(matchedToken.sentIndex() != sentINdx || matchedToken.index() < reprMent.startIndex || matchedToken.index() > reprMent.endIndex ){
                for(int i = reprMent.startIndex; i<reprMent.endIndex; i++){
                    CoreLabel matchedLabel = corefSentenceTokens.get(i-1); //resolved.add(tokens.get(i).word());
                    repTokens.add(matchedLabel);
                }
            }
            else{
                //repTokens.add( matchedToken );
            }
        }
        return repTokens;
    }


    public Map<Integer,Map<Integer,Set<String>>> getSentenceEndtokenEntitySet(String text, Annotation document){
        text += ".";
        EntityLinkingResult result = this.textEntityLinker.getTextEntityLinkingResult(text, document);
        System.out.println("\n\n[Entity Linking and Coreference Result]---------------\n\n");
        System.out.println(result.toString());
        System.out.println("\n\n");
        return result.sentenceTokenEntitySet;
    }

    public Set<String> getEntitiesOfToken(CoreLabel matchedToken){
        Integer sentNum = matchedToken.sentIndex() + 1;
        Integer tokenIndex = matchedToken.index();
        Set<String> entities = null;
        if(this.sentenceEndtokenEntitySet.containsKey(sentNum)){
            if(this.sentenceEndtokenEntitySet.get(sentNum).containsKey(tokenIndex)){
                entities = this.sentenceEndtokenEntitySet.get(sentNum).get(tokenIndex);
            }
        }
        return entities;
    }




    public void printContexts(List<ContextResult> contextResults){
        for(ContextResult contextResult: contextResults){
            System.out.println("\n----------------------------");
            System.out.println("[context = " + contextResult.context + " ]");
            System.out.println("[entities = {");
            for(String entity: contextResult.entites){
                System.out.println(entity);
            }
            System.out.println("}]");
        }
    }








    /****************************************************************
     * MAIN
     */
    public static void main(String[] args){
        String text = "The most popular products of Apple, which is a computer producer, include Iphone and Macbook, however its Watch is a niche product.";
        if(args.length > 0){
            text = "";
            for(String arg:args)
                text += " " + arg;
        }
        System.out.println("[text = " + text + "]\n\n\n");
        ContextAnnotation contextAnnotation = new ContextAnnotation();
        List<ContextResult> contexts = contextAnnotation.getContextsOfDocument( text );

        contextAnnotation.printContexts( contexts );

        System.out.println("\n\nParse Finished!\n\n");
    }

}
